/*
 * @Author: 姚潘涛
 * @Date: 2024-06-03 10:41:29
 * @LastEditors: 姚潘涛
 * @LastEditTime: 2024-10-10 19:17:00
 * @Description:
 *
 * Copyright (c) 2024 by pandaman, All Rights Reserved.
 */
#include "xict_common/perception/detection/yolo_detect.h"

namespace xict_common {
    namespace perception {
        namespace detection {
            /**
             * @brief Construct a new Yolo Detect:: Yolo Detect object
             *
             * @param model_path 模型路径
             * @param use_CUDA 是否使用cuda加速
             */
            YoloDetect::YoloDetect(std::string model_path, bool use_CUDA) {
                if (yolo.ReadModel(net, model_path, use_CUDA)) {
                    LOG_F(INFO, "read net ok!");
                } else {
                    LOG_F(INFO, "read net failed!");
                    return;
                }
            }

            YoloDetect::~YoloDetect() { }

            /**
             * @brief yolo识别函数
             *
             * @param img 待检测图像
             * @return cv::Mat 识别之后的图像
             */
            cv::Mat YoloDetect::Detect(cv::Mat img) {
                cv::Mat result_img;
                // 生成随机颜色
                srand(time(0));
                for (int i = 0; i < 80; i++) {
                    int b = rand() % 256;
                    int g = rand() % 256;
                    int r = rand() % 256;
                    color.push_back(cv::Scalar(b, g, r));
                }

                if (yolo.Detect(img, net, result)) {
                    result_img = yolo.DrawPred(img, result, color);
                    return result_img;
                } else {
                    LOG_F(INFO, "Detect Failed!");
                    return img;
                }
            }

            /**
             * @brief 识别的结果，包含bbox数据
             *
             * @param img 待检测图像
             * @return struct struct_yolo_result 识别后的图像及bbox数据
             */
            struct struct_yolo_result YoloDetect::DetectBbox(cv::Mat img) {
                cv::Mat result_img;
                struct struct_yolo_result result_;
                // 生成随机颜色
                srand(time(0));
                for (int i = 0; i < 80; i++) {
                    int b = rand() % 256;
                    int g = rand() % 256;
                    int r = rand() % 256;
                    color.push_back(cv::Scalar(b, g, r));
                }

                if (yolo.Detect(img, net, result)) {
                    result_img     = yolo.DrawPred(img, result, color);
                    result_.img    = result_img;
                    result_.result = result;
                    return result_;
                    // return result_img;
                } else {
                    LOG_F(INFO, "Detect Failed!");
                    // return img;
                    result_.img = img;
                    return result_;
                }
            }
        }   // namespace detection
    }       // namespace perception
}   // namespace xict_common
